import os
import re
import pickle
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import sklearn
from sklearn.model_selection import train_test_split
from sklearn import metrics
import skimage
import skimage.color
import skimage.io
import skimage.feature
import skimage.transform
from sklearn.base import BaseEstimator, TransformerMixin
from skimage.color import rgb2gray
from sklearn.linear_model import SGDClassifier
from sklearn.model_selection import GridSearchCV
import sklearn.metrics
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV
from mathe import rgb2gray_transform, hogtransformer

# load the data
data = pickle.load(open('./pickle_files/data_animals_head_20.pickle','rb'))
print("Load pickle data")

dataDesc = data['description']  # There are 20 classes and 2057 images are there. All the images are 80 x 80 (rgb)
X = data['data']                # Contains all images, 
y = data['target']              # Ground truth
labels = data['labels']         # Label names of all dataset

# split the data into train and test
x_train,x_test,y_train,y_test = train_test_split(X,y,test_size=0.2,stratify=y)
print("x_train.shape: {}; x_test.shape: {}; len(y_train): {}; len(y_test): {}".format(x_train.shape,x_test.shape,len(y_train),len(y_test)))

# After we get this best model, we also get the best parameters like followed:
# Best parameter is:  {'hogtransform__cells_per_block': (2, 2), 'hogtransform__orientations': 8, 
# 'hogtransform__pixels_per_cell': (8, 8), 'sgd__learning_rate': 'optimal', 'sgd__loss': 'hinge'}

from sklearn.pipeline import make_pipeline
pipeline1 = make_pipeline(rgb2gray_transform(),
            hogtransformer(orientations=8,
                            pixels_per_cell=(8,8),
                            cells_per_block=(3,3)))
feature_vector = pipeline1.fit_transform(x_train)
# standard scaler
scalar = StandardScaler()
transformed_xtrain = scalar.fit_transform(feature_vector)
model = SGDClassifier(learning_rate='optimal',loss='hinge',alpha=0.01,early_stopping=True)
model.fit(transformed_xtrain,y_train)
# evaluate
feature_vector = pipeline1.fit_transform(x_test)
transformed_x = scalar.transform(feature_vector)
y_pred_test = model.predict(transformed_x)
cr = sklearn.metrics.classification_report(y_test,y_pred_test,output_dict=True)
print(pd.DataFrame(cr).T)
print("Model evaluation score: ", metrics.cohen_kappa_score(y_test,y_pred_test))

# save models for flask app
pickle.dump(model,open('./pickle_files/dsa_image_classification_sgd.pickle','wb'))

pickle.dump(scalar,open('./pickle_files/dsa_scaler.pickle','wb'))